package com.github.cyf.interceptor.tenant;

import com.github.cyf.util.MSUtil;
import com.github.cyf.util.PrintUtil;
import com.github.cyf.context.SecurityUtil;
import lombok.NoArgsConstructor;
import lombok.Setter;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import java.lang.reflect.Method;
import java.sql.Connection;

/**
 * 为什么拦截 StatementHandler 的 prepare 方法 ？
 * 因为我们不止查询语句需要，还有更新、删除语句也需要，就不像前面一样拦截 Executors 的 query 方法就可以了
 * Executors 在具体执行查询，修改前都会调用 StatementHandler 的 prepare 方法进行预编译
 * https://blog.csdn.net/qq_30631063/article/details/118861569
 * @author chenyifan
 * @create 2024-04-15 15:29
 */
@Intercepts(
        @Signature(
                type = StatementHandler.class,
                method = "prepare",
                args = {Connection.class, Integer.class}
        )
)
@Setter
@NoArgsConstructor
public class TenantInterceptor implements Interceptor {

    private static PrintUtil printUtil = new PrintUtil(TenantInterceptor.class);

    /**
     * 数据库租户id字段名称
     */
    private String tenantIdColumn = "tenant_id";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        printUtil.printInterceptorBegin();

        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
        MappedStatement ms = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
        Method method = MSUtil.getMapperMethod(ms);
        // 被 IgnoreTenant 注解的方法，不进行租户id筛选
        if (method != null && method.getAnnotation(IgnoreTenant.class) != null) {
            return invocation.proceed();
        }
        // 修改 sql，添加筛选条件
        String sql = (String) metaObject.getValue("delegate.boundSql.sql");
        String tenantSql = addTenantFilter(sql, SecurityUtil.getTenantId());
        printUtil.printSql(ms.getId(), tenantSql);
        metaObject.setValue("delegate.boundSql.sql", tenantSql);

        printUtil.printInterceptorEnd();
        // 执行目标方法
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    /**
     * 根据原始 sql，添加租户id筛选条件，只能对简单 sql 做修改。
     * @param originalSQL
     * @param tenantId
     * @return
     */
    private String addTenantFilter(String originalSQL, Long tenantId) {
        originalSQL = originalSQL.toUpperCase();
        if (originalSQL.contains("INSERT")) {
            return originalSQL;
        }
        // 检查原始 SQL 是否已经有 WHERE 子句
        int whereIndex = originalSQL.indexOf("WHERE");
        // 检查原始 SQL 是否已经有 ORDER BY 子句
        int orderByIndex = originalSQL.indexOf("ORDER BY");
        // 检查原始 SQL 是否已经有 LIMIT 子句
        int limiIndex = originalSQL.indexOf("LIMIT");

        StringBuilder conBuilder = new StringBuilder(originalSQL.length() + 15);
        // 构建筛选条件部分
        if (whereIndex != -1) {
            // 如果原始 SQL 已经有 WHERE 子句，使用 AND 连接租户筛选条件
            conBuilder.append(" AND ");
        } else {
            // 原始 SQL 没有 WHERE 子句，直接添加 WHERE 子句和租户筛选条件
            conBuilder.append(" WHERE ");
        }
        conBuilder.append(tenantIdColumn).append(" = ").append(tenantId).append(" ");

        // 寻找插入位置
        int insertIndex;
        if (orderByIndex != -1) {
            insertIndex = orderByIndex;
        } else if (limiIndex != -1) {
            insertIndex = limiIndex;
        } else {
            insertIndex = originalSQL.length();
        }
        StringBuilder sqlBuilder = new StringBuilder(originalSQL);
        sqlBuilder.insert(insertIndex, conBuilder);
        return sqlBuilder.toString();
    }
}
